/**
 * WebSocket client for ComfyUI real-time events
 */

import type {
  ComfyUIEvent,
  EventListener,
  EventListeners,
  WSServerMessage
} from '../types';

// Support both browser and Node.js environments
let WebSocketImpl: typeof WebSocket;
if (typeof WebSocket !== 'undefined') {
  // Browser environment
  WebSocketImpl = WebSocket;
} else {
  // Node.js environment
  try {
    const ws = require('ws');
    WebSocketImpl = ws.default || ws;
  } catch (error) {
    throw new Error('WebSocket implementation not found. Please install the "ws" package for Node.js environments.');
  }
}

/**
 * WebSocket connection states
 */
export enum ConnectionState {
  DISCONNECTED = 'disconnected',
  CONNECTING = 'connecting',
  CONNECTED = 'connected',
  RECONNECTING = 'reconnecting',
  FAILED = 'failed'
}

/**
 * WebSocket client configuration
 */
export interface WebSocketClientConfig {
  url: string;
  clientId?: string;
  reconnectAttempts?: number;
  reconnectDelay?: number;
  heartbeatInterval?: number;
}

/**
 * WebSocket client for real-time ComfyUI events
 */
export class WebSocketClient {
  private ws: WebSocket | null = null;
  private readonly config: Required<WebSocketClientConfig>;
  private state: ConnectionState = ConnectionState.DISCONNECTED;
  private reconnectCount = 0;
  private heartbeatTimer: NodeJS.Timeout | null = null;
  private reconnectTimer: NodeJS.Timeout | null = null;
  
  private readonly eventListeners = new Map<string, Set<EventListener>>();

  /**
   * Creates a new WebSocketClient instance
   * @param config - WebSocket configuration
   */
  constructor(config: WebSocketClientConfig) {
    this.config = {
      url: config.url,
      clientId: config.clientId || this.generateClientId(),
      reconnectAttempts: config.reconnectAttempts || 5,
      reconnectDelay: config.reconnectDelay || 1000,
      heartbeatInterval: config.heartbeatInterval || 30000
    };
  }

  /**
   * Gets the current connection state
   */
  get connectionState(): ConnectionState {
    return this.state;
  }

  /**
   * Gets the client ID
   */
  get clientId(): string {
    return this.config.clientId;
  }

  /**
   * Checks if the client is connected
   */
  get isConnected(): boolean {
    return this.state === ConnectionState.CONNECTED && this.ws?.readyState === WebSocket.OPEN;
  }

  /**
   * Connects to the WebSocket server
   * @returns Promise that resolves when connected
   */
  async connect(): Promise<void> {
    if (this.state === ConnectionState.CONNECTING || this.isConnected) {
      return;
    }

    return new Promise((resolve, reject) => {
      try {
        this.setState(ConnectionState.CONNECTING);
        
        const wsUrl = `${this.config.url}?clientId=${this.config.clientId}`;
        this.ws = new WebSocketImpl(wsUrl);

        this.ws.onopen = () => {
          this.setState(ConnectionState.CONNECTED);
          this.reconnectCount = 0;
          this.startHeartbeat();
          this.emit({
            type: 'connected',
            timestamp: new Date(),
            clientId: this.config.clientId
          });
          resolve();
        };

        this.ws.onmessage = (event) => {
          this.handleMessage(event);
        };

        this.ws.onclose = (event) => {
          this.handleClose(event);
        };

        this.ws.onerror = (error) => {
          this.handleError(error);
          if (this.state === ConnectionState.CONNECTING) {
            reject(new Error('Failed to connect to WebSocket'));
          }
        };

      } catch (error) {
        this.setState(ConnectionState.FAILED);
        reject(error);
      }
    });
  }

  /**
   * Disconnects from the WebSocket server
   */
  disconnect(): void {
    this.stopHeartbeat();
    this.stopReconnect();
    
    if (this.ws) {
      this.ws.close(1000, 'Client disconnect');
      this.ws = null;
    }
    
    this.setState(ConnectionState.DISCONNECTED);
    this.emit({
      type: 'disconnected',
      timestamp: new Date(),
      reason: 'Client disconnect'
    });
  }

  /**
   * Adds an event listener
   * @param eventType - Event type to listen for
   * @param listener - Event listener function
   */
  on<T extends ComfyUIEvent>(eventType: T['type'], listener: EventListener<T>): void {
    if (!this.eventListeners.has(eventType)) {
      this.eventListeners.set(eventType, new Set());
    }
    this.eventListeners.get(eventType)!.add(listener as EventListener);
  }

  /**
   * Removes an event listener
   * @param eventType - Event type
   * @param listener - Event listener function to remove
   */
  off<T extends ComfyUIEvent>(eventType: T['type'], listener: EventListener<T>): void {
    const listeners = this.eventListeners.get(eventType);
    if (listeners) {
      listeners.delete(listener as EventListener);
      if (listeners.size === 0) {
        this.eventListeners.delete(eventType);
      }
    }
  }

  /**
   * Adds multiple event listeners at once
   * @param listeners - Object with event listeners
   */
  addListeners(listeners: Partial<EventListeners>): void {
    if (listeners.progress) this.on('progress', listeners.progress);
    if (listeners.executing) this.on('executing', listeners.executing);
    if (listeners.executed) this.on('executed', listeners.executed);
    if (listeners.error) this.on('error', listeners.error);
    if (listeners.connected) this.on('connected', listeners.connected);
    if (listeners.disconnected) this.on('disconnected', listeners.disconnected);
    if (listeners.queued) this.on('queued', listeners.queued);
    if (listeners.started) this.on('started', listeners.started);
    if (listeners.completed) this.on('completed', listeners.completed);
  }

  /**
   * Removes all event listeners for a specific event type
   * @param eventType - Event type to clear (optional, clears all if not provided)
   */
  removeAllListeners(eventType?: string): void {
    if (eventType) {
      this.eventListeners.delete(eventType);
    } else {
      this.eventListeners.clear();
    }
  }

  /**
   * Handles incoming WebSocket messages
   * @private
   * @param event - WebSocket message event
   */
  private handleMessage(event: MessageEvent): void {
    try {
      const message: WSServerMessage = JSON.parse(event.data);
      this.processServerMessage(message);
    } catch (error) {
      console.error('Failed to parse WebSocket message:', error);
    }
  }

  /**
   * Processes server messages and emits appropriate events
   * @private
   * @param message - Server message
   */
  private processServerMessage(message: WSServerMessage): void {
    const timestamp = new Date();

    switch (message.type) {
      case 'progress': {
        const progressData = (message as any).data;
        this.emit({
          type: 'progress',
          timestamp,
          promptId: progressData.prompt_id || '',
          nodeId: progressData.node || '',
          progress: progressData.value || 0,
          max: progressData.max || 1
        });
        break;
      }

      case 'executing': {
        const executingData = (message as any).data;
        if (executingData.node) {
          this.emit({
            type: 'executing',
            timestamp,
            promptId: executingData.prompt_id || '',
            nodeId: executingData.node
          });
        }
        break;
      }

      case 'executed': {
        const executedData = (message as any).data;
        this.emit({
          type: 'executed',
          timestamp,
          promptId: executedData.prompt_id || '',
          nodeId: executedData.node || '',
          outputs: executedData.output || {},
          executionTime: 0 // Server doesn't provide this
        });
        break;
      }

      case 'execution_start': {
        const startData = (message as any).data;
        this.emit({
          type: 'started',
          timestamp,
          promptId: startData.prompt_id || ''
        });
        break;
      }

      case 'execution_success': {
        const successData = (message as any).data;
        this.emit({
          type: 'completed',
          timestamp,
          promptId: successData.prompt_id || '',
          outputs: {},
          executionTime: 0
        });
        break;
      }

      case 'execution_error': {
        const errorData = (message as any).data;
        this.emit({
          type: 'error',
          timestamp,
          promptId: errorData.prompt_id || '',
          nodeId: errorData.node_id,
          error: new Error(errorData.exception_message || 'Unknown error'),
          details: {
            nodeType: errorData.node_type,
            executed: errorData.executed,
            exceptionType: errorData.exception_type,
            traceback: errorData.traceback,
            currentInputs: errorData.current_inputs,
            currentOutputs: errorData.current_outputs
          }
        });
        break;
      }

      case 'status': {
        // Status messages contain queue information and system status
        // We can emit these as status events for monitoring
        const statusData = (message as any).data;
        this.emit({
          type: 'status',
          timestamp,
          data: statusData
        });
        break;
      }

      case 'execution_cached': {
        // Execution was cached, no need to run again
        const cachedData = (message as any).data;
        this.emit({
          type: 'cached',
          timestamp,
          promptId: cachedData.prompt_id || '',
          nodes: cachedData.nodes || []
        });
        break;
      }

      default:
        // Only log truly unknown message types, not the ones we handle above
        if (!['status', 'execution_cached'].includes(message.type)) {
          console.log('Unknown message type:', message.type);
        }
    }
  }

  /**
   * Handles WebSocket close events
   * @private
   * @param event - Close event
   */
  private handleClose(event: CloseEvent): void {
    this.stopHeartbeat();
    
    if (this.state !== ConnectionState.DISCONNECTED) {
      this.emit({
        type: 'disconnected',
        timestamp: new Date(),
        reason: `Connection closed: ${event.code} ${event.reason}`
      });

      // Attempt reconnection if not manually disconnected
      if (event.code !== 1000 && this.reconnectCount < this.config.reconnectAttempts) {
        this.attemptReconnect();
      } else {
        this.setState(ConnectionState.FAILED);
      }
    }
  }

  /**
   * Handles WebSocket errors
   * @private
   * @param error - Error event
   */
  private handleError(error: Event): void {
    console.error('WebSocket error:', error);
    this.emit({
      type: 'error',
      timestamp: new Date(),
      error: new Error('WebSocket connection error')
    });
  }

  /**
   * Attempts to reconnect to the server
   * @private
   */
  private attemptReconnect(): void {
    if (this.state === ConnectionState.RECONNECTING) {
      return;
    }

    this.setState(ConnectionState.RECONNECTING);
    this.reconnectCount++;

    const delay = this.config.reconnectDelay * Math.pow(2, this.reconnectCount - 1);
    
    this.reconnectTimer = setTimeout(async () => {
      try {
        await this.connect();
      } catch (error) {
        if (this.reconnectCount < this.config.reconnectAttempts) {
          this.attemptReconnect();
        } else {
          this.setState(ConnectionState.FAILED);
        }
      }
    }, delay);
  }

  /**
   * Starts the heartbeat timer
   * @private
   */
  private startHeartbeat(): void {
    this.heartbeatTimer = setInterval(() => {
      if (this.isConnected) {
        this.ws?.send(JSON.stringify({ type: 'ping' }));
      }
    }, this.config.heartbeatInterval);
  }

  /**
   * Stops the heartbeat timer
   * @private
   */
  private stopHeartbeat(): void {
    if (this.heartbeatTimer) {
      clearInterval(this.heartbeatTimer);
      this.heartbeatTimer = null;
    }
  }

  /**
   * Stops the reconnect timer
   * @private
   */
  private stopReconnect(): void {
    if (this.reconnectTimer) {
      clearTimeout(this.reconnectTimer);
      this.reconnectTimer = null;
    }
  }

  /**
   * Sets the connection state
   * @private
   * @param state - New connection state
   */
  private setState(state: ConnectionState): void {
    this.state = state;
  }

  /**
   * Emits an event to all listeners
   * @private
   * @param event - Event to emit
   */
  private emit(event: ComfyUIEvent): void {
    const listeners = this.eventListeners.get(event.type);
    if (listeners) {
      for (const listener of listeners) {
        try {
          listener(event);
        } catch (error) {
          console.error('Error in event listener:', error);
        }
      }
    }
  }

  /**
   * Generates a unique client ID
   * @private
   * @returns Unique client ID
   */
  private generateClientId(): string {
    return `client_${Date.now()}_${Math.random().toString(36).substring(2, 8)}`;
  }
}
